from yolov11trt.config import *
import threading
import math
import time

# ======================
# PID控制器类（保持不变）
# ======================
class PID:
    def __init__(self, P, I, D):
        self.Kp = P
        self.Ki = I
        self.Kd = D
        self.previous_error = 0
        self.integral = 0

    def compute(self, set_point, measured_value):
        error = measured_value - set_point
        self.integral += error
        derivative = error - self.previous_error
        output = self.Kp * error + self.Ki * self.integral + self.Kd * derivative
        self.previous_error = error
        return output

    def update_pid(self, P, I, D):
        self.Kp, self.Ki, self.Kd = P, I, D

    def get_p(self): return self.Kp
    def get_i(self): return self.Ki
    def get_d(self): return self.Kd
    
    def set_p(self, P): self.Kp = P
    def set_i(self, I): self.Ki = I
    def set_d(self, D): self.Kd = D

# ======================
# 模拟云台控制器（打印版）
# ======================
class ServoController:
    def __init__(self):
        # 初始化角度为中位
        self.pan_angle = 90
        self.tilt_angle = 90

    @staticmethod
    def degrees_to_radians(degrees):
        """角度转弧度简洁实现"""
        return math.radians(degrees)

    def set_angle(self, pan_angle, tilt_angle):
        """设置云台角度（角度制），自动转换为弧度并打印"""
        # 角度限制保护
        pan_angle = max(0, min(180, pan_angle))
        tilt_angle = max(0, min(180, tilt_angle))
        
        # 转换为弧度
        pan_rad = self.degrees_to_radians(pan_angle)
        tilt_rad = self.degrees_to_radians(tilt_angle)
        
        # 更新当前角度
        self.pan_angle = pan_angle
        self.tilt_angle = tilt_angle
        
        # 打印输出（格式：PAN_ANG,TILT_ANG,PAN_RAD,TILT_RAD）
        print(f"云台角度设置:")
        print(f"水平角度: {pan_angle:.2f}° ({pan_rad:.6f}rad)")
        print(f"俯仰角度: {tilt_angle:.2f}° ({tilt_rad:.6f}rad)")
        print("-" * 40)

# ======================
# 云台追踪系统（保持不变）
# ======================
class YuntaiTrack:
    def __init__(self):
        self.servo_controller = ServoController()
        self.pan_angle = 90
        self.tilt_angle = 90
        self.pan_pid = PID(0.022, 0.0, 0.013)
        self.tilt_pid = PID(0.022, 0.0, 0.013)

    def track(self, target_x, target_y, img_width, img_height):
        # 计算目标与图像中心的偏差
        error_x = target_x - (img_width / 2)
        error_y = target_y - (img_height / 2)

        # PID算法计算舵机调整角度
        pan_adjust = self.pan_pid.compute(0, error_x)
        tilt_adjust = self.tilt_pid.compute(0, error_y)

        # 更新舵机角度
        new_pan_angle = self.pan_angle + pan_adjust
        new_tilt_angle = self.tilt_angle - tilt_adjust
        
        # 限制角度范围
        new_pan_angle = max(min(new_pan_angle, 180), 0)
        new_tilt_angle = max(min(new_tilt_angle, 180), 60)

        # 更新角度并打印输出
        self.pan_angle = new_pan_angle
        self.tilt_angle = new_tilt_angle
        self.servo_controller.set_angle(self.pan_angle, self.tilt_angle)

# ======================
# 云台控制线程（修改版）
# ======================
class YuntaiThread(threading.Thread):
    def __init__(self, yuntai_track, img_width, img_height):
        super().__init__()
        self.yuntai_track = yuntai_track
        self.img_width = img_width
        self.img_height = img_height
        self.target_x = None
        self.target_y = None
        self.target_name = None
        self.lock = threading.Lock()
        self.condition = threading.Condition()
        self.stop_event = threading.Event()

    def set_target_name(self, target_name):
        with self.lock:
            self.target_name = target_name

    def handle_infer_result(self, infer_result):
        if self.target_name is None:
            return
        
        boxs = infer_result.get("boxs", [])
        classid = infer_result.get("classid", [])
        
        if boxs is None:
            return
        
        for i in range(len(boxs)):
            if categories[int(classid[i])] == self.target_name:
                x1, y1, x2, y2 = boxs[i]
                self.target_x = (x1 + x2) / 2
                self.target_y = (y1 + y2) / 2
                with self.condition:
                    self.condition.notify()
                break
        else:
            self.target_x = None
            self.target_y = None

    def run(self):
        while not self.stop_event.is_set():
            with self.condition:
                self.condition.wait()
            if self.target_x is not None and self.target_y is not None:
                self.yuntai_track.track(
                    self.target_x, 
                    self.target_y, 
                    self.img_width, 
                    self.img_height
                )

    def stop(self):
        self.stop_event.set()
        with self.condition:
            self.condition.notify()

